{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e484e6a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | default_exp _components.encoder.json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "330229f3",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/kumaran/.local/lib/python3.11/site-packages/pydantic/_internal/_config.py:257: UserWarning: Valid config keys have changed in V2:\n",
      "* 'json_encoders' has been removed\n",
      "  warnings.warn(message, UserWarning)\n"
     ]
    }
   ],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "import json\n",
    "from typing import *\n",
    "\n",
    "from pydantic import BaseModel\n",
    "\n",
    "from fastkafka._components.logger import get_logger\n",
    "from fastkafka._components.meta import export"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84ed548a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import tempfile\n",
    "\n",
    "from pydantic import Field\n",
    "\n",
    "from fastkafka._components.logger import suppress_timestamps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84c22900",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40c6f220",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "logger = get_logger(__name__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e61e48a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INFO] __main__: ok\n"
     ]
    }
   ],
   "source": [
    "suppress_timestamps()\n",
    "logger = get_logger(__name__, level=20)\n",
    "logger.info(\"ok\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44803f0e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "def _to_json_utf8(o: Any) -> bytes:\n",
    "    \"\"\"Converts to JSON and then encodes with UTF-8\"\"\"\n",
    "    if hasattr(o, \"model_dump_json\"):\n",
    "        return o.model_dump_json().encode(\"utf-8\")  # type: ignore\n",
    "    else:\n",
    "        return json.dumps(o).encode(\"utf-8\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6adf56fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert _to_json_utf8({\"a\": 1, \"b\": [2, 3]}) == b'{\"a\": 1, \"b\": [2, 3]}'\n",
    "\n",
    "\n",
    "class A(BaseModel):\n",
    "    name: str = Field()\n",
    "    age: int\n",
    "\n",
    "\n",
    "assert _to_json_utf8(A(name=\"Davor\", age=12)) == b'{\"name\":\"Davor\",\"age\":12}'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7fb9e74d",
   "metadata": {},
   "outputs": [],
   "source": [
    "class User(BaseModel):\n",
    "    name: str\n",
    "    favorite_number: Optional[int] = None\n",
    "    favorite_color: Optional[str] = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "659da394",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "@export(\"fastkafka.encoder\")\n",
    "def json_encoder(msg: BaseModel) -> bytes:\n",
    "    \"\"\"\n",
    "    Encoder to encode pydantic instances to json string\n",
    "\n",
    "    Args:\n",
    "        msg: An instance of pydantic basemodel\n",
    "\n",
    "    Returns:\n",
    "        Json string in bytes which is encoded from pydantic basemodel\n",
    "    \"\"\"\n",
    "    return _to_json_utf8(msg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a89f790",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "b'{\"name\":\"Kumaran\",\"favorite_number\":9,\"favorite_color\":\"black\"}'"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "msg = User(name=\"Kumaran\", favorite_number=9, favorite_color=\"black\")\n",
    "\n",
    "\n",
    "actual = json_encoder(msg)\n",
    "display(actual)\n",
    "\n",
    "assert isinstance(actual, bytes)\n",
    "assert actual == b'{\"name\":\"Kumaran\",\"favorite_number\":9,\"favorite_color\":\"black\"}'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cdbdbb80",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "@export(\"fastkafka.encoder\")\n",
    "def json_decoder(raw_msg: bytes, cls: Type[BaseModel]) -> Any:\n",
    "    \"\"\"\n",
    "    Decoder to decode json string in bytes to pydantic model instance\n",
    "\n",
    "    Args:\n",
    "        raw_msg: Bytes message received from Kafka topic\n",
    "        cls: Pydantic class; This pydantic class will be used to construct instance of same class\n",
    "\n",
    "    Returns:\n",
    "        An instance of given pydantic class\n",
    "    \"\"\"\n",
    "    msg_dict = json.loads(raw_msg.decode(\"utf-8\"))\n",
    "\n",
    "    return cls(**msg_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c98ae71a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "User(name='123', favorite_number=0, favorite_color='111')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "raw_msg = b'{\"name\": \"123\", \"favorite_number\": 0, \"favorite_color\": \"111\"}'\n",
    "\n",
    "\n",
    "actual = json_decoder(raw_msg, cls=User)\n",
    "display(actual)\n",
    "\n",
    "assert isinstance(actual, User)\n",
    "assert actual.name == \"123\"\n",
    "assert actual.favorite_number == 0\n",
    "assert actual.favorite_color == \"111\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9cf574fc",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
